#coding:utf-8

import numpy as np
import itertools
from scipy.optimize import minimize
from AssetAllocation.Common import *
from AssetAllocation.BasicOptimize import BasicOptimize


# risk budgeting optimization
def calculate_portfolio_var(w, V):
    # function that calculates portfolio risk
    w = np.matrix(w)
    return (w * V * w.T)[0, 0]


# MRC 边际风险贡献 RC单个资产风险贡献
def calculate_risk_contribution(w, V):
    # function that calculates asset contribution to total risk
    w = np.matrix(w)
    sigma = np.sqrt(calculate_portfolio_var(w, V))
    # Marginal Risk Contribution
    MRC = V * w.T / sigma
    # Risk Contribution
    RC = np.multiply(MRC, w.T)
    return RC


# calculate portfolio risk
def calc_portfolio_risk_volatility(x, pars):
    V = pars[0]  # covariance table
    x_t = pars[1]  # risk target in percent of portfolio risk
    sig_p = np.sqrt(calculate_portfolio_var(x, V))  # portfolio sigma
    risk_target = np.asmatrix(np.multiply(sig_p, x_t))
    asset_RC = calculate_risk_contribution(x, V)
    J = sum(np.square(asset_RC - risk_target.T))[0, 0]  # sum of squared error

    # 笛卡尔
    # res = [x for x in itertools.product(asset_RC,repeat = 2)]
    # res = [x[0] - x[1] for x in res]
    # J = sum(np.square(res))
    return J


def calc_portfolio_risk_mdd(x, pars):
    return None


def calc_portfolio_risk_downside_risk(x, pars):
    return None


def total_weight_constraint(x):
    return np.sum(x) - 1.0


def long_only_constraint(x):
    return x


class RiskBudget(BasicOptimize):
    def __init__(self, symbol_list=[], annualized_factor=252):
        super().__init__()
        self.symbol_list = symbol_list
        self.num_assets = len(self.symbol_list)
        self.annualized_factor = annualized_factor
        self.annualized_covariance_matrix = None
        # self.annualized_volatility = None
        self.mdd = None
        self.annualized_downside_risk = None

    def initialize_with_returns(self, df_returns):
        self.symbol_list = df_returns.columns.tolist()
        self.annualized_covariance_matrix = (df_returns.cov() * self.annualized_factor).values
        a = 0
        # self.mdd = xxx
        # self.annualized_downside_risk = xxx

    def initialize_with_params(self, correlation, expected_volatility_array):
        self.annualized_covariance_matrix = Calc_Covariance_Matrix(correlation, expected_volatility_array)

    def initialize_with_params_covariance(self, annualized_covariance_matrix):
        self.annualized_covariance_matrix = annualized_covariance_matrix

    # 根据资产预期目标风险贡献度来计算各资产的权重
    # risk_budget = [0.25, 0.25, 0.25, 0.25]
    # 意思是目标是让四个资产风险贡献度相等，即都为25%
    def optimize(self, optimize_objective, risk_budget=[]):

        # 确定资产数量
        num_assets = 0
        optimize_objective_function = None
        args = []
        if optimize_objective == "VOLATILITY":
            covariance = self.annualized_covariance_matrix
            num_assets = np.array(covariance).shape[0]
            optimize_objective_function = calc_portfolio_risk_volatility
            args = [covariance, risk_budget]
        #
        elif optimize_objective == "MDD":
            optimize_objective_function = calc_portfolio_risk_mdd
            num_assets = len(self.mdd)
            args = [self.mdd, risk_budget]
        #
        elif optimize_objective == "DOWNSIDERISK":
            optimize_objective_function = calc_portfolio_risk_downside_risk()
            # num_assets = len(self.annualized_downside_risk)
            # args = [self.annualized_downside_risk, risk_budget]

        #
        if len(risk_budget) != num_assets:
            print("risk budget length error")
            return None

        # 初始化资产权重
        w0 = np.ones(num_assets) * 1.0 / num_assets
        #     x_t = w0
        cons = [{'type': 'eq', 'fun': total_weight_constraint},
                {'type': 'ineq', 'fun': long_only_constraint},
                # {'type': 'eq', 'fun': lambda w: w[0] - 0.2},  # 活期
                # {'type': 'eq', 'fun': lambda w: w[1] - 0.2},  # 债券
                # {'type': 'eq', 'fun': lambda w: w[3] - 0.1},  # 黄金
                ]
        # cons = [{'type': 'eq', 'fun': lambda x: np.sum(x) - 1}]
        bnds, instrument_type_cons = self._process_constraints()
        cons += instrument_type_cons

        sol = minimize(optimize_objective_function, w0,
                       args=args,
                       method='SLSQP',
                       bounds=bnds,
                       constraints=cons,
                       options={'disp': False, 'ftol': 1e-15})
        res = []
        for i in range(len(sol.x)):
            if sol.x[i] < 1e-6:
                res.append(0.0)
            else:
                res.append(round(sol.x[i], 4))
        return np.array(res)

    def get_target_risk_portfolio(self, target_risk):
        pass


# 随机序列，测试
def Test_With_Random_Data():
    #
    num_assets = 4
    num_obs = 100
    df_returns = Generate_Random_Data(num_assets, num_obs)
    print(df_returns)

    rb = RiskBudget()
    rb.initialize_with_returns(df_returns)
    #
    # rb.set_symbol_constraints([(0.5, 1), (0.3, 1), (0,1), (0,1)])
    risk_budget = [0.2, 0.2, 0.2, 0.2]
    weights = rb.optimize("VOLATILITY", risk_budget)
    print(weights)

    # rb.optimize("MDD", risk_budget)


def Test_With_Params():
    #
    annual_std_dev = pd.Series([0.20, 0.30, 0.02, 0.038])
    correlation = [[1.000000, 0.8404, -0.0256, -0.0065],
                   [0, 1.000000, -0.0377, 0.0020],
                   [0, 0, 1.000000, 0.2268],
                   [0, 0, 0, 1.000000]]
    #
    risk_budget = [0.25, 0.25, 0.25, 0.25]
    #
    rb = RiskBudget()
    rb.initialize_with_params(correlation, annual_std_dev)
    #
    weights = rb.optimize("VOLATILITY", risk_budget)
    print(weights)


def Test_With_Params_Covariance():
    #
    covariance =[[1.3614494176132513e-08, 1.8334838053826713e-08, 6.349686760814494e-08, -7.016927888405647e-09, 1.9748037204597332e-08],
                 [1.8334838053826713e-08, 5.124438204356057e-07, 2.389458264258538e-07, 1.1413829128050989e-07, 1.3570292626308216e-07],
                 [6.349686760814494e-08, 2.389458264258538e-07, 0.0001529740765588536, 1.722700847131028e-06, 1.0627824149621483e-05],
                 [-7.016927888405647e-09, 1.1413829128050989e-07, 1.722700847131028e-06, 5.521884764008903e-05, -4.052578161170029e-07],
                 [1.9748037204597332e-08, 1.3570292626308216e-07, 1.0627824149621483e-05, -4.052578161170029e-07, 4.1904661880923344e-05]]

    # 目标风险占比，即各占20%，平均风险
    risk_budget = [0.2, 0.2, 0.2, 0.2, 0.2]

    rb = RiskBudget(symbol_list=[])
    rb.initialize_with_params_covariance(covariance)
    weights = rb.optimize("VOLATILITY", risk_budget)
    print(weights)

    # [0.2, 0.2, 0.1593, 0.1, 0.3407] 2020-9-11
    # [0.8281, 0.1337, 0.008, 0.0146, 0.0156] 2020-9-11 屏蔽掉活期，债券，黄金
    # [0.4397, 0.4567, 0.0218, 0.0394, 0.0424]


if __name__ == '__main__':

    Test_With_Params()
    # Test_With_Params_Covariance()
    # Test_With_Random_Data()
